import torch
import torch.nn.functional
from torch.utils.data import Dataset
import torch.nn.functional as F


class BaseDataset(Dataset):
    def __init__(self, 
                 data: torch.Tensor, 
                 graph: torch.Tensor, 
                 seq: int, 
                 pred_len: int, 
                 target_turbine: int, 
                 target_feature: int, 
                 base_feature: int, 
                 logfile=None, 
                 single_point=True
                 ):
        '''
        创建数据集：需要每个点的特征data，静态图graph，和必要的参数

        Args:
            data (torch.Tensor): 每个点的特征，一般直接从DataFrame转过来得到
            graph (torch.Tensor): N × N的静态图
            seq (int): 回望窗口长度
            pred_len (int): 预测长度
            target_turbine (int): 预测目标风机编号
            target_feature (int): 预测目标特征编号，一般都是差分
            base_feature (int): 真正的预测目标特征，因为一般都是差分，所以这个变量用于后续预测结果的还原
            logfile (str): log所在位置
            signle_point (bool): 是否为单点预测
        '''
        super(BaseDataset, self).__init__()
        self.data = data
        self.graph = graph
        self.t, self.n, self.c = data.shape
        self.seq = seq
        self.pred_len = pred_len
        self.target_turbine = target_turbine
        self.target_feature = target_feature
        self.base_feature = base_feature
        self.logfile = logfile
        self.signle_point = single_point


    def __len__(self):
        '''
        一个样本的长度为seq + pred_len，因此总共的样本数需要减去它再+1 
        '''
        return len(self.data) - self.seq - self.pred_len + 1
    

    def __getitem__(self, index):
        """
        第index个样本包含了：
        1.x: 一个包括了回望窗口的序列，包括当前时间点向前seq个点上所有风机的所有特征，尺寸为(seq, N, C)
        2.graph: 相似度矩阵，尺寸为(N, N)
        3.target: 目标风机在目标特征上的值
        4.base: 目标风机在前一个点基础特征上的值
        
        Args:
            index (int) :样本编号
        """
        x = self.getdata(index)
        future_start = index + self.seq
        if self.signle_point:
            target = self.data[future_start, self.target_turbine, self.target_feature]
            base = self.data[future_start - 1, self.target_turbine, self.base_feature]
        else:
            target = self.data[future_start: future_start + self.pred_len, self.target_turbine, self.target_feature]
            base = self.data[future_start - 1: future_start + self.pred_len - 1, self.target_turbine, self.base_feature]
        return x, self.graph, target, base


    def getdata(self, index):
        pass


    def split(self, split_rate):
        """
        将数据集按照一定的比例拆分，用于划分训练集，验证集和测试集
        
        Args: 
            split_rate (list(float)): 划分比例，相加为1的浮点数列表
        """
        length = len(self.data) - self.seq - self.pred_len + 1
        start = 0
        end = 0
        res = []
        for rate in split_rate:
            end += rate
            # 按照比例转换子数据集的起点和终点下标，并切分出相应的特征数据和时间相似度数据
            data = self.data[int(start * length): int(end * length) + self.seq + self.pred_len - 1]
            res.append(self.__class__(data, self.graph, self.seq, self.pred_len, self.target_turbine, self.target_feature, self.base_feature))
            start = end
        return res